package com.seafood.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket消息处理器
 * 
 * @author System
 * @since 2024-01-01
 */
@Slf4j
@Component
public class WebSocketHandler extends TextWebSocketHandler {

    private final ObjectMapper objectMapper = new ObjectMapper();
    
    // 存储所有的WebSocket会话
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    
    // 存储用户ID和Session的映射关系
    private final Map<Long, String> userSessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket连接建立，sessionId: {}", session.getId());
        
        sessions.put(session.getId(), session);
        
        // 从URI中获取用户ID参数（如果有的话）
        String userId = getUserIdFromSession(session);
        if (userId != null) {
            try {
                Long userIdLong = Long.parseLong(userId);
                userSessions.put(userIdLong, session.getId());
                log.info("用户 {} 连接WebSocket，sessionId: {}", userIdLong, session.getId());
            } catch (NumberFormatException e) {
                log.warn("无效的用户ID: {}", userId);
            }
        }
        
        // 发送连接成功消息
        sendMessage(session, new WebSocketMessage("system", "connection", "连接成功", null));
        
        log.info("当前WebSocket连接数: {}", sessions.size());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        log.info("收到WebSocket消息，sessionId: {}, 消息内容: {}", session.getId(), message.getPayload());
        
        try {
            // 解析消息
            WebSocketMessage wsMessage = objectMapper.readValue(message.getPayload(), WebSocketMessage.class);
            
            // 处理消息
            handleMessage(session, wsMessage);
            
        } catch (Exception e) {
            log.error("处理WebSocket消息失败", e);
            sendMessage(session, new WebSocketMessage("system", "error", "消息处理失败", null));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误，sessionId: {}", session.getId(), exception);
        
        // 清理会话
        cleanupSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接关闭，sessionId: {}, 关闭状态: {}", session.getId(), closeStatus);
        
        // 清理会话
        cleanupSession(session);
        
        log.info("当前WebSocket连接数: {}", sessions.size());
    }

    /**
     * 处理消息
     */
    private void handleMessage(WebSocketSession session, WebSocketMessage message) throws IOException {
        String type = message.getType();
        
        switch (type) {
            case "ping":
                // 心跳检测
                sendMessage(session, new WebSocketMessage("system", "pong", "心跳响应", null));
                break;
                
            case "subscribe":
                // 订阅通知
                handleSubscribe(session, message);
                break;
                
            case "unsubscribe":
                // 取消订阅
                handleUnsubscribe(session, message);
                break;
                
            default:
                log.warn("未知的消息类型: {}", type);
                sendMessage(session, new WebSocketMessage("system", "error", "未知的消息类型", null));
        }
    }

    /**
     * 处理订阅
     */
    private void handleSubscribe(WebSocketSession session, WebSocketMessage message) throws IOException {
        String topic = (String) message.getData();
        log.info("用户订阅主题: {}, sessionId: {}", topic, session.getId());
        
        // 在session中存储订阅信息
        session.getAttributes().put("subscribed_" + topic, true);
        
        sendMessage(session, new WebSocketMessage("system", "subscribed", "订阅成功: " + topic, topic));
    }

    /**
     * 处理取消订阅
     */
    private void handleUnsubscribe(WebSocketSession session, WebSocketMessage message) throws IOException {
        String topic = (String) message.getData();
        log.info("用户取消订阅主题: {}, sessionId: {}", topic, session.getId());
        
        // 从session中移除订阅信息
        session.getAttributes().remove("subscribed_" + topic);
        
        sendMessage(session, new WebSocketMessage("system", "unsubscribed", "取消订阅成功: " + topic, topic));
    }

    /**
     * 向指定用户发送消息
     */
    public void sendMessageToUser(Long userId, WebSocketMessage message) {
        String sessionId = userSessions.get(userId);
        if (sessionId != null) {
            WebSocketSession session = sessions.get(sessionId);
            if (session != null && session.isOpen()) {
                try {
                    sendMessage(session, message);
                    log.info("向用户 {} 发送消息成功", userId);
                } catch (IOException e) {
                    log.error("向用户 {} 发送消息失败", userId, e);
                }
            } else {
                log.warn("用户 {} 的会话不存在或已关闭", userId);
                userSessions.remove(userId);
            }
        } else {
            log.warn("用户 {} 未连接WebSocket", userId);
        }
    }

    /**
     * 广播消息到所有订阅指定主题的客户端
     */
    public void broadcastToTopic(String topic, WebSocketMessage message) {
        log.info("广播消息到主题: {}, 消息: {}", topic, message.getContent());
        
        int sentCount = 0;
        for (WebSocketSession session : sessions.values()) {
            if (session.isOpen() && isSubscribedToTopic(session, topic)) {
                try {
                    sendMessage(session, message);
                    sentCount++;
                } catch (IOException e) {
                    log.error("广播消息失败，sessionId: {}", session.getId(), e);
                }
            }
        }
        
        log.info("广播消息完成，主题: {}, 发送数量: {}", topic, sentCount);
    }

    /**
     * 广播消息到所有连接的客户端
     */
    public void broadcastToAll(WebSocketMessage message) {
        log.info("广播消息到所有客户端，消息: {}", message.getContent());
        
        int sentCount = 0;
        for (WebSocketSession session : sessions.values()) {
            if (session.isOpen()) {
                try {
                    sendMessage(session, message);
                    sentCount++;
                } catch (IOException e) {
                    log.error("广播消息失败，sessionId: {}", session.getId(), e);
                }
            }
        }
        
        log.info("广播消息完成，发送数量: {}", sentCount);
    }

    /**
     * 发送消息到指定会话
     */
    private void sendMessage(WebSocketSession session, WebSocketMessage message) throws IOException {
        if (session.isOpen()) {
            String messageJson = objectMapper.writeValueAsString(message);
            session.sendMessage(new TextMessage(messageJson));
        }
    }

    /**
     * 检查会话是否订阅了指定主题
     */
    private boolean isSubscribedToTopic(WebSocketSession session, String topic) {
        return session.getAttributes().containsKey("subscribed_" + topic);
    }

    /**
     * 从会话中获取用户ID
     */
    private String getUserIdFromSession(WebSocketSession session) {
        String query = session.getUri().getQuery();
        if (query != null && query.contains("userId=")) {
            String[] params = query.split("&");
            for (String param : params) {
                if (param.startsWith("userId=")) {
                    return param.substring(7);
                }
            }
        }
        return null;
    }

    /**
     * 清理会话
     */
    private void cleanupSession(WebSocketSession session) {
        sessions.remove(session.getId());
        
        // 清理用户会话映射
        userSessions.entrySet().removeIf(entry -> entry.getValue().equals(session.getId()));
    }

    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return sessions.size();
    }

    /**
     * 获取在线用户数
     */
    public int getOnlineUserCount() {
        return userSessions.size();
    }
}